import numpy as np
import math
from numpy.core.defchararray import multiply
from numpy.core.fromnumeric import shape
import PIL.Image as Image

def gause1DKernel(kernel_size, sigma):
    """
    1D Gause Kernel.
    """
    SUM = 0.0
    kernel = np.zeros(shape = (1, kernel_size))
    if kernel_size % 2 == 0:
        center = kernel_size // 2 - 0.5
    else:
        center = kernel_size // 2

    for i in range(kernel_size):
        x = i - center
        kernel[0][i] = math.exp( - (x**2) / (2 * sigma**2))
        SUM += kernel[0][i]
    kernel = kernel / SUM
    return kernel

def gauseKernel(kernel_size, sigma):
    """
    2D Gause Kernel.
    """
    SUM = 0.0
    kernel = np.zeros(shape = (kernel_size, kernel_size))
    if kernel_size % 2 == 0:
        center = kernel_size // 2 - 0.5
    else:
        center = kernel_size // 2
    for i in range(kernel_size):
        for j in range(kernel_size):
            x = i - center
            y = j - center
            kernel[i][j] = math.exp( - (x**2 + y**2) / (2 * sigma**2))
            SUM = SUM + kernel[i][j]
    
    kernel = kernel / SUM
    return kernel

def BilateralKernel(kernel_size, sigma1, sigma2, image):
    """
    2D Bilateral Kernel.
    """
    
    inputs = np.array(image).astype("float64")
    kernel = np.zeros(shape = (inputs.shape[0], inputs.shape[1], kernel_size, kernel_size))
    print(kernel.shape)
    center = kernel_size // 2
    for i in range(inputs.shape[0]):
        for j in range(inputs.shape[1]):
            SUM = 0.0
            for m in range(kernel_size):
                for n in range(kernel_size):
                    x = m - center
                    y = n - center
                    if i + x >= 0 and i + x < inputs.shape[0] and j + y >= 0 and j + y < inputs.shape[1]:
                        kernel[i][j][m][n] = math.exp( - (x**2 + y**2) / (2 * sigma1 ** 2) - (inputs[i][j] - inputs[i + x][j + y]) ** 2 / (2 * sigma2 ** 2))
                    else:
                        kernel[i][j][m][n] = 0.0
                    SUM = SUM + kernel[i][j][m][n]
    
            kernel[i][j] = kernel[i][j] / SUM
    return kernel

def BilateralFilter(kernel_size, sigma1, sigma2, image):
    inputs = np.array(image)
    outputs = np.zeros(shape=(inputs.shape[0], inputs.shape[1]))
    kernel = BilateralKernel(kernel_size, sigma1, sigma2, inputs)
    center = kernel_size // 2
    for i in range(inputs.shape[0]):
        for j in range(inputs.shape[1]):
            tmp = 0.0
            for m in range(kernel_size):
                for n in range(kernel_size):
                    x = m - center
                    y = n - center
                    if i + x >= 0 and i + x < inputs.shape[0] and j + y >= 0 and j + y < inputs.shape[1]:
                        tmp += kernel[i][j][m][n] * inputs[i + x][j + y]
            # print(tmp)
            outputs[i][j] = tmp
    outputImage = Image.fromarray(outputs.astype("uint")).convert("L")
    outputImage.save("./GK/BilateralFiltered.jpg")

class Conv2D():
    """
    A very simple implement of convolutional layer.
    """

    def __init__(self, filters, kernel_size, sigma, boundary = "zero", padding = "valid"):
        self.filters = [gauseKernel(kernel_size, (sigma + i)) for i in range(filters)]
        self.kernel_size = kernel_size
        self.sigma = sigma
        self.padding = padding
        self.outputs = []
        self.boundary = boundary
    
    def subOfTwoKernels(self, inputs, kernel1, kernel2):
        self.filters = [kernel1, kernel2]
        outputs = self.call(inputs)
        return outputs[0] - outputs[1]

    def boundaryChange(self, inputs):
        image = np.array(inputs)
        h = image.shape[0]
        w = image.shape[1]
        
        imageExtend = np.zeros(shape = (3*h,3*w))
        for i in range(h, 2*h):
            for j in range(w, 2*w):
                imageExtend[i][j] = image[i-h][j-w]
        if self.boundary == "wrap":
            for i in range(3*h):
                for j in range(3*w):
                    imageExtend[i][j] = image[i % h][j % w]
        elif self.boundary == "copy":
            for i in range(h):
                for j in range(w,2*w):
                    imageExtend[i][j] = imageExtend[h][j]
            for i in range(2*h, 3*h):
                for j in range(w,2*w):
                    imageExtend[i][j] = imageExtend[2*h-1][j]
            for i in range(h,2*h):
                for j in range(w):
                    imageExtend[i][j] = imageExtend[i][w]
            for i in range(h,2*h):
                for j in range(2*w,3*w):
                    imageExtend[i][j] = imageExtend[i][2*w-1]
            for i in range(h):
                for j in range(w):
                    imageExtend[i][j] = imageExtend[h][w]
            for i in range(h):
                for j in range(2*w, 3*w):
                    imageExtend[i][j] = imageExtend[h][2*w-1]
            for i in range(2*h,3*h):
                for j in range(w):
                    imageExtend[i][j] = imageExtend[2*h-1][w]
            for i in range(2*h,3*h):
                for j in range(2*w,3*w):
                    imageExtend[i][j] = imageExtend[2*h-1][2*w-1]
        outputImage = Image.fromarray(imageExtend.astype("uint")).convert("L")
        name = "./GK/rawImageWith" + self.boundary + ".jpg"
        outputImage.save(name)
        return imageExtend
                    

    def call(self, inputs):
        inputs = self.boundaryChange(inputs)
        image = np.array(inputs)
        h = len(image) // 3
        w = len(image[0]) // 3
        if self.padding == "SAME" or self.padding == "same":
            for f in range(len(self.filters)):
                output = np.zeros(shape = (h, w))
                for i in range(h):
                    for j in range(w):
                        output[i][j] = self.multiply_element_wide(i, j, h, w, image, self.filters[f])
                self.outputs.append(output)
        elif self.padding == "VALID" or self.padding == "valid":
            for f in range(len(self.filters)):
                k = self.kernel_size // 2
                output = np.zeros(shape = (h - 2*k, w - 2*k))
                for i in range(k, h - k):
                    for j in range(k, w - k):
                        output[i - k][j - k] = self.multiply_element_wide(i, j, h, w, image, self.filters[f])
                self.outputs.append(output)
        else:
            raise Exception("PaddingError, please enter valid or same.")
        return self.outputs

    def multiply_element_wide(self, i, j, h, w, image, kernel):
        SUM = 0.0
        k = self.kernel_size // 2
        center = k
        for m in range(-k, k+1):
            for n in range(-k, k+1):
                # if i + m >= 0 and j + n >= 0 and i + m < h and j + n < w:
                tmp = kernel[center + m][center + n] * image[h + i + m][w + j + n]
                SUM += tmp
        return SUM
    
    def __call__(self, inputs, **kwds):
        return self.call(inputs)

class Conv1D():
    """
    A very simple implement of convolutional layer.
    """

    def __init__(self, filters, kernel_size, stride, padding = "valid"):
        self.filters = [gause1DKernel(kernel_size, 1) for i in range(filters)]
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.outputs = []
    
    def call(self, inputs):
        image = np.array(inputs)
        h = len(image)
        w = len(image[0])
        if self.padding == "SAME" or self.padding == "same":
            for f in range(len(self.filters)):
                output = np.zeros(shape = (h, w))
                for i in range(h):
                    for j in range(w):
                        output[i][j] = self.multiply_element_wide(i, j, h, w, image, self.filters[f])
                self.outputs.append(output)
        elif self.padding == "VALID" or self.padding == "valid":
            for f in range(len(self.filters)):
                k = self.kernel_size // 2
                output = np.zeros(shape = (h - 2*k, w - 2*k))
                for i in range(k, h - k):
                    for j in range(k, w - k):
                        output[i - k][j - k] = self.multiply_element_wide(i, j, h, w, image, self.filters[f])
                self.outputs.append(output)
        else:
            raise Exception("PaddingError, please enter valid or same.")
        return self.outputs

    def multiply_element_wide(self, i, j, h, w, image, kernel):
        SUM = 0.0
        k = self.kernel_size // 2
        center = k
        for n in range(-k, k+1):
            if j + n >= 0 and j + n < w:
                tmp = kernel[0][center + n] * image[i][j + n]
                SUM += tmp
        return SUM
    
    def __call__(self, inputs, **kwds):
        return self.call(inputs)


# inputs = [[1,1,1,1,1],
#           [1,1,1,1,1],
#           [1,1,1,1,1],
#           [1,1,1,1,1],
#           [1,1,1,1,1]]

# inputs = [[1,51,1,1,1],
#           [1,17,15,1,1],
#           [1,11,2,17,1],
#           [1,23,1,1,1],
#           [1,1,1,1,65]]

# c = Conv2D(1, 3, 3, "SAME")
# outputs = c(inputs)
# print(outputs[0])
# kernel1 = gauseKernel(3,1000000)
# kernel2 = gauseKernel(3,1)
# cc = Conv2D(1,3,1,"SAME")
# output = cc.subOfTwoKernels(inputs, kernel1, kernel2)
# print(output)
# output = cc(inputs)
# imageD = np.array(inputs) - output[0]
# print(imageD)
# print(gauseKernel(3,1000000))
# print(gauseKernel(3,1))
# print(np.matmul(gause1DKernel(3,1).T, gause1DKernel(3,1)))
# print(gause1DKernel(3,1))
# cc = Conv1D(1, 3, 1, "SAME")
# outputs1D = cc(inputs)
# print(outputs1D)

def single2DKernelFiltering(kernel_size, sigma, boundary, padding):
    """
    一个2D高斯核，图片是1.jpg，进行滤波
    """
    imageRGB = Image.open("1.jpg")
    imageL = imageRGB.convert("L")
    imageL.save("./GK/rawImage.jpg")
    imageL_arr = np.array(imageL)
    conv = Conv2D(1, kernel_size, sigma, boundary, padding)
    output = conv(imageL_arr)

    outputImage = Image.fromarray(output[0].astype("int")).convert("L")
    name = "./GK/2DfilteredImage" + str(sigma) + boundary + padding + ".jpg"
    outputImage.save(name)
    return output[0]

def single1DKernelFiltering(kernel_size, sigma, padding):
    """
    一个1D高斯核，图片是1.jpg，进行滤波
    """
    imageRGB = Image.open("1.jpg")
    imageL = imageRGB.convert("L")
    imageL.save("./GK/rawImage.jpg")
    imageL_arr = np.array(imageL)
    conv = Conv1D(1, kernel_size, sigma, padding)
    output = conv(imageL_arr)

    outputImage = Image.fromarray(output[0].astype("int")).convert("L")
    name = "./GK/1DfilteredImage" + str(sigma) + padding + ".jpg"
    outputImage.save(name)
    return output[0]

def detailExtract():
    """
    通过一个高斯滤波和原图相减，提取图像细节
    """
    imageRGB = Image.open("1.jpg")
    imageL = imageRGB.convert("L")
    imageL_arr = np.array(imageL)
    conv = Conv2D(1, 3, 1, "zero", "SAME")
    output = conv(imageL_arr)
    afterFiltering_arr = output[0]
    detail = imageL_arr - afterFiltering_arr
    outputImage = Image.fromarray(detail.astype("uint32")).convert("L")
    outputImage.save("./GK/detailExtracted.jpg")
    return detail


def sharpen(r):
    """
    通过提取的细节进行锐化
    """
    detail = detailExtract()
    # print(detail)
    imageRGB = Image.open("1.jpg")
    imageL = imageRGB.convert("L")
    # imageL.show()
    imageL_arr = np.array(imageL).astype("float64")
    # print(imageL_arr)
    imageL_arr += r * detail
    # print(imageL_arr.astype("uint32"))
    outputImage = Image.fromarray(imageL_arr.astype("uint32")).convert("L")
    outputImage.save("./GK/shapenedImage.jpg")
    return imageL_arr

def subOfTwoKernel(size1, sigma1, size2, sigma2):
    imageRGB = Image.open("1.jpg")
    imageL = imageRGB.convert("L")
    # imageL.show()
    imageL_arr = np.array(imageL).astype("float64")
    kernel1 = gauseKernel(size1, sigma1)
    kernel2 = gauseKernel(size2, sigma2)
    cc = Conv2D(1, 3, 3, "zero", "SAME")
    output = cc.subOfTwoKernels(imageL_arr, kernel1, kernel2)
    subImage = Image.fromarray(output.astype("uint")).convert("L")
    subImage.save("./GK/subImage.jpg")
    # print(output)


def generate2DKernelByTwo1DKernels():
    print(np.matmul(gause1DKernel(3,1).T, gause1DKernel(3,1)))

def generate2DWith1Ds(kernel_size, sigma):
    kernel2D = (np.matmul(gause1DKernel(kernel_size, sigma).T, gause1DKernel(kernel_size, sigma)))
    return kernel2D

two = single2DKernelFiltering(3, 1, "zero", "SAME")
single2DKernelFiltering(3, 1, "copy", "SAME")
single2DKernelFiltering(3, 1, "zero", "valid")
single2DKernelFiltering(3, 1, "wrap", "valid")
single2DKernelFiltering(3, 7, "zero", "SAME")
single2DKernelFiltering(3, 7, "copy", "SAME")
single2DKernelFiltering(3, 7, "zero", "valid")
single2DKernelFiltering(3, 7, "wrap", "valid")
one = single2DKernelFiltering(3, 100, "zero", "SAME")
print(two - one)
single2DKernelFiltering(3, 100, "copy", "SAME")
single2DKernelFiltering(3, 100, "zero", "valid")
single2DKernelFiltering(3, 100, "wrap", "valid")
single1DKernelFiltering(3, 1, "SAME")
single1DKernelFiltering(3, 1, "valid")
single1DKernelFiltering(3, 3, "SAME")
single1DKernelFiltering(3, 3, "valid")
# detailExtract()
# sharpen(0.5)
# subOfTwoKernel(3, 1, 3, 100000)
# generatedKernel = generate2DWith1Ds(3, 1)
# print(generatedKernel)
# image1 = Image.open("./GK/rawImage.jpg")
# BilateralFilter(3, 1, 1, image1)